import cv2
import numpy as np
from typing import Optional, Tuple, Dict, Any
import time
import threading
from queue import Queue, Empty, Full
from concurrent.futures import ThreadPoolExecutor, as_completed
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
# 添加项目根目录到Python路径
sys.path.append(os.path.join(current_dir, '../'))  # 根据实际位置调整层级

# 假设这些模块已经存在
from image_provider import ImageProvider
from logger import setup_logger
from seg_inference2 import seg_func
from pose_inference import pose_func
from py_utils.rknn_executor import RKNN_model_container

# 初始化日志
logger = setup_logger()

class ImageData:
    """图像数据结构"""
    def __init__(self):
        self.timestamp: float = 0.0
        self.readable_timestamp: str = ""
        self.color_img: np.ndarray = None
        self.depth_img: np.ndarray = None
        self.color_path: Optional[str] = None

class RKNNPoolExecutor:
    """RKNN线程池执行器"""
    
    def __init__(self, rknn_model_path: str, inference_func, tpes: int = 1, npus: list = [0]):
        self.tpes = tpes
        self.inference_func = inference_func
        self.npus = npus
        self.model_index = 0
        
        # 初始化RKNN模型池
        logger.info(f"初始化RKNN模型池: {rknn_model_path}, TPEs={tpes}, NPUs={npus}")
        self.rknn_pool = self._init_rknn_models(rknn_model_path, tpes, npus)
        
        # 创建线程池
        total_workers = tpes * len(npus)
        self.pool = ThreadPoolExecutor(max_workers=total_workers)
        
        # 任务管理
        self.task_id = 0
        self.pending_futures = {}
        self.completed_results = {}
        self.lock = threading.Lock()
        
        logger.info(f"RKNN线程池初始化完成，总工作线程数: {total_workers}")
    
    def _init_rknn_model(self, model_path: str, npu_id: int = 0):
        """初始化单个RKNN模型"""
        rknn_lite = RKNN_model_container(model_path, "rk3588", npu_id)
        return rknn_lite
    
    def _init_rknn_models(self, model_path: str, tpes: int, npus: list):
        """初始化多个RKNN模型"""
        rknn_list = []
        for npu in npus:
            for i in range(tpes):
                rknn_list.append(self._init_rknn_model(model_path, npu))
        return rknn_list
    
    def put(self, frame: np.ndarray) -> int:
        """提交推理任务并返回任务ID"""
        with self.lock:
            task_id = self.task_id
            self.task_id += 1
        
        # 选择RKNN模型（轮询）
        model_index = self.model_index % len(self.rknn_pool)
        self.model_index += 1
        
        logger.debug(f"提交任务 {task_id}: 使用模型索引={model_index}")
        
        # 包装推理函数
        def inference_wrapper(model, img):
            try:
                start_time = time.time()
                result = self.inference_func(model, img)
                end_time = time.time()
                logger.debug(f"推理任务 {task_id} 完成，耗时: {(end_time-start_time)*1000:.1f}ms")
                return result
            except Exception as e:
                logger.error(f"推理任务 {task_id} 异常: {e}")
                return None, False
        
        # 提交任务
        future = self.pool.submit(inference_wrapper, self.rknn_pool[model_index], frame)
        
        # 保存任务
        with self.lock:
            self.pending_futures[task_id] = future
        
        # 添加完成回调
        future.add_done_callback(lambda f, tid=task_id: self._task_completed(tid, f))
        
        return task_id
    
    def _task_completed(self, task_id: int, future):
        """任务完成回调"""
        try:
            result = future.result()
            success = result is not None and result[1] if isinstance(result, tuple) else result is not None
        except Exception as e:
            logger.error(f"任务 {task_id} 执行失败: {e}")
            result = None
            success = False
        
        with self.lock:
            # 移动到完成队列
            self.completed_results[task_id] = (result, success)
            # 从待处理队列移除
            if task_id in self.pending_futures:
                del self.pending_futures[task_id]
    
    def get(self):
        """获取最早完成的结果"""
        with self.lock:
            if len(self.completed_results) == 0:
                return None, None, False
            
            # 获取最早的任务ID
            earliest_task_id = min(self.completed_results.keys())
            result, success = self.completed_results.pop(earliest_task_id)
            return earliest_task_id, result, success
    
    def has_results(self):
        """检查是否有结果可获取"""
        with self.lock:
            return len(self.completed_results) > 0
    
    def pending_count(self):
        """获取待处理任务数量"""
        with self.lock:
            return len(self.pending_futures)
    
    def close(self):
        """关闭线程池"""
        logger.info("关闭RKNN线程池...")
        self.pool.shutdown(wait=True)
        for rknn_lite in self.rknn_pool:
            rknn_lite.release()
        logger.info("RKNN线程池已关闭")

class SegPoseRK:
    """分割和姿态估计异步处理器"""
    
    def __init__(self, 
                 seg_model_path: str = "./models/merge_seg_int8_rk3588.rknn",
                 pose_model_path: str = "./models/merge_pose_int8_rk3588.rknn",
                 seg_workers: int = 6,
                 pose_workers: int = 6,
                 seg_npus: list = [1, 2],
                 pose_npus: list = [0],
                 max_queue_size: int = 50,
                 target_fps: int = 30):
        
        self.seg_model_path = seg_model_path
        self.pose_model_path = pose_model_path
        self.max_queue_size = max_queue_size
        self.target_fps = target_fps
        self.running = False
        
        # 初始化RKNN线程池
        logger.info("初始化分割和姿态估计线程池...")
        self.seg_pool = RKNNPoolExecutor(
            rknn_model_path=seg_model_path,
            inference_func=seg_func,
            tpes=seg_workers,
            npus=seg_npus
        )
        
        self.pose_pool = RKNNPoolExecutor(
            rknn_model_path=pose_model_path,
            inference_func=pose_func,
            tpes=pose_workers,
            npus=pose_npus
        )
        
        # 任务管理
        self.task_queue = Queue(maxsize=max_queue_size)
        self.result_queue = Queue(maxsize=max_queue_size)
        self.pending_tasks: Dict[int, Dict] = {}
        self.next_task_id = 0
        self.lock = threading.Lock()
        
        # 性能统计
        self.submit_count = 0
        self.process_count = 0
        self.total_inference_time = 0
        self.total_frames_processed = 0
        self.last_log_time = time.time()
        
        # 线程引用
        self.threads = []
    
    def start(self):
        """启动异步处理器"""
        if self.running:
            logger.warning("处理器已经在运行中")
            return
        
        self.running = True
        
        # 启动任务分发线程
        dispatch_thread = threading.Thread(target=self._dispatch_tasks, name="TaskDispatcher", daemon=True)
        dispatch_thread.start()
        self.threads.append(dispatch_thread)
        
        # 启动结果收集线程
        collector_thread = threading.Thread(target=self._result_collector, name="ResultCollector", daemon=True)
        collector_thread.start()
        self.threads.append(collector_thread)
        
        # 启动性能监控线程
        monitor_thread = threading.Thread(target=self._performance_monitor, name="PerformanceMonitor", daemon=True)
        monitor_thread.start()
        self.threads.append(monitor_thread)
        
        logger.info("SegPose异步处理器启动完成")
    
    def stop(self):
        """停止异步处理器"""
        if not self.running:
            return
        
        logger.info("正在停止SegPose处理器...")
        self.running = False
        
        # 等待线程结束
        for thread in self.threads:
            thread.join(timeout=2.0)
            if thread.is_alive():
                logger.warning(f"线程 {thread.name} 未能及时结束")
        
        # 清理线程池
        self.seg_pool.close()
        self.pose_pool.close()
        
        logger.info("SegPose处理器已停止")
    
    def put(self, image_data: ImageData) -> int:
        """提交处理任务"""
        if not self.running:
            logger.error("处理器未启动，无法提交任务")
            return -1
        
        with self.lock:
            task_id = self.next_task_id
            self.next_task_id += 1
        
        # 创建任务
        task = {
            'id': task_id,
            'image_data': image_data,
            'seg_id': None,
            'pose_id': None,
            'seg_result': None,
            'pose_result': None,
            'seg_success': False,
            'pose_success': False,
            'start_time': time.time(),
            'state': 'queued'
        }
        
        # 保存任务
        with self.lock:
            self.pending_tasks[task_id] = task
        
        # 放入任务队列
        try:
            self.task_queue.put(task, timeout=0.1)
            self.submit_count += 1
            logger.debug(f"任务 {task_id} 已提交")
            return task_id
        except Full:
            logger.warning(f"任务队列已满，任务 {task_id} 被丢弃")
            with self.lock:
                del self.pending_tasks[task_id]
            return -1
    
    def _dispatch_tasks(self):
        """任务分发线程"""
        while self.running or not self.task_queue.empty():
            try:
                task = self.task_queue.get(timeout=0.5)
                
                # 提交到seg和pose线程池
                seg_id = self.seg_pool.put(task['image_data'].color_img)
                pose_id = self.pose_pool.put(task['image_data'].color_img)
                
                # 更新任务状态
                with self.lock:
                    if task['id'] in self.pending_tasks:
                        task['seg_id'] = seg_id
                        task['pose_id'] = pose_id
                        task['state'] = 'processing'
                        logger.debug(f"任务 {task['id']} 已分发: seg_id={seg_id}, pose_id={pose_id}")
                
            except Empty:
                continue
            except Exception as e:
                logger.error(f"任务分发错误: {e}")
    
    def _result_collector(self):
        """结果收集线程"""
        check_interval = 0.05  # 50ms检查一次
        
        while self.running or len(self.pending_tasks) > 0:
            results_found = False
            
            # 检查seg结果
            seg_id, seg_result, seg_success = self.seg_pool.get()
            if seg_id is not None:
                task = self._find_task_by_seg_id(seg_id)
                if task:
                    with self.lock:
                        task['seg_result'] = seg_result
                        task['seg_success'] = seg_success
                    results_found = True
                    logger.debug(f"收到seg结果: task_id={task['id']}")
            
            # 检查pose结果
            pose_id, pose_result, pose_success = self.pose_pool.get()
            if pose_id is not None:
                task = self._find_task_by_pose_id(pose_id)
                if task:
                    with self.lock:
                        task['pose_result'] = pose_result
                        task['pose_success'] = pose_success
                    results_found = True
                    logger.debug(f"收到pose结果: task_id={task['id']}")
            
            # 处理完成的任务
            self._process_completed_tasks()
            
            # 如果没有结果，短暂休眠
            if not results_found:
                time.sleep(check_interval)
    
    def _find_task_by_seg_id(self, seg_id: int):
        """通过seg_id查找任务"""
        with self.lock:
            for task in self.pending_tasks.values():
                if task.get('seg_id') == seg_id:
                    return task
        return None
    
    def _find_task_by_pose_id(self, pose_id: int):
        """通过pose_id查找任务"""
        with self.lock:
            for task in self.pending_tasks.values():
                if task.get('pose_id') == pose_id:
                    return task
        return None
    
    def _process_completed_tasks(self):
        """处理已完成的任务"""
        completed_task_ids = []
        
        with self.lock:
            for task_id, task in list(self.pending_tasks.items()):
                # 检查是否两个结果都完成
                if task['seg_success'] and task['pose_success']:
                    # 计算推理时间
                    inference_time = time.time() - task['start_time']
                    self.total_inference_time += inference_time
                    self.total_frames_processed += 1
                    
                    # 构建结果
                    result = {
                        'task_id': task_id,
                        'image_data': task['image_data'],
                        'seg_result': task['seg_result'],
                        'pose_result': task['pose_result'],
                        'seg_success': task['seg_success'],
                        'pose_success': task['pose_success'],
                        'inference_time': inference_time,
                    }
                    
                    # 放入结果队列
                    try:
                        self.result_queue.put(result, timeout=0.1)
                        completed_task_ids.append(task_id)
                        self.process_count += 1
                        logger.debug(f"任务 {task_id} 完成，推理时间: {inference_time*1000:.1f}ms")
                    except Full:
                        logger.warning("结果队列已满，丢弃结果")
                    
                    # 清理任务
                    del self.pending_tasks[task_id]
        
        return len(completed_task_ids) > 0
    
    def _performance_monitor(self):
        """性能监控线程"""
        while self.running:
            time.sleep(2.0)  # 每2秒输出一次统计
            self._log_performance()
    
    def _log_performance(self):
        """记录性能统计"""
        current_time = time.time()
        elapsed = current_time - self.last_log_time
        
        if elapsed >= 2.0:
            # 计算帧率
            submit_fps = self.submit_count / elapsed if elapsed > 0 else 0
            process_fps = self.process_count / elapsed if elapsed > 0 else 0
            
            # 计算平均推理时间
            avg_inference = 0
            if self.total_frames_processed > 0:
                avg_inference = self.total_inference_time / self.total_frames_processed
            
            # 队列状态
            with self.lock:
                pending_tasks = len(self.pending_tasks)
            
            seg_pending = self.seg_pool.pending_count()
            pose_pending = self.pose_pool.pending_count()
            
            logger.info(
                f"性能统计: 提交={submit_fps:.1f}FPS, "
                f"处理={process_fps:.1f}FPS, "
                f"队列深度={pending_tasks}, "
                f"seg待处理={seg_pending}, "
                f"pose待处理={pose_pending}, "
                f"平均延迟={avg_inference*1000:.1f}ms"
            )
            
            # 重置计数器
            self.submit_count = 0
            self.process_count = 0
            self.last_log_time = current_time
    
    def get_result(self, block: bool = False, timeout: Optional[float] = None):
        """获取处理结果"""
        try:
            return self.result_queue.get(block=block, timeout=timeout)
        except Empty:
            return None
    
    def has_result(self) -> bool:
        """检查是否有结果可获取"""
        return not self.result_queue.empty()
    
    def pending_count(self) -> int:
        """获取待处理任务数量"""
        with self.lock:
            return len(self.pending_tasks)
    
    def get_status(self) -> Dict[str, Any]:
        """获取处理器状态"""
        with self.lock:
            pending_tasks = len(self.pending_tasks)
        
        return {
            'running': self.running,
            'pending_tasks': pending_tasks,
            'seg_pending': self.seg_pool.pending_count(),
            'pose_pending': self.pose_pool.pending_count(),
            'result_queue_size': self.result_queue.qsize(),
            'total_processed': self.total_frames_processed,
            'avg_inference_time': (
                self.total_inference_time / self.total_frames_processed 
                if self.total_frames_processed > 0 else 0
            )
        }

# 使用示例
def main():
    """使用示例"""
    # 创建SegPose处理器
    seg_pose = SegPoseRK(
        seg_model_path="./models/merge_yolov11_seg.rknn",
        pose_model_path="./models/merge_yolov11_pose.rknn",
        seg_workers=6,
        pose_workers=6,
        seg_npus=[1, 2],
        pose_npus=[0],
        target_fps=30
    )
    
    # 启动处理器
    seg_pose.start()
    
    try:
        # 模拟图像处理
        image_provider = ImageProvider(from_folder=True)
        
        # 提交任务
        for i in range(100):
            image_data = image_provider.get_next_image()
            if image_data is None:
                break
            
            task_id = seg_pose.put(image_data)
            if task_id == -1:
                logger.warning(f"任务 {i} 提交失败")
            
            # 获取结果
            while seg_pose.has_result():
                result = seg_pose.get_result()
                print(result)
                if result:
                    logger.info(f"处理完成: 任务ID={result['task_id']}, "
                               f"推理时间={result['inference_time']*1000:.1f}ms")
            
            time.sleep(1.0 / 30)  # 30FPS
        
        # 等待所有任务完成
        while seg_pose.pending_count() > 0:
            result = seg_pose.get_result(block=True, timeout=1.0)
            if result:
                logger.info(f"最终结果: 任务ID={result['task_id']}")
    
    except KeyboardInterrupt:
        logger.info("接收到中断信号")
    
    finally:
        # 停止处理器
        seg_pose.stop()
        logger.info("程序结束")

if __name__ == "__main__":
    main()